// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: jhkim

#include <cxxopts.hpp>

#include "coin/proto/coin_market_enums.pb.h"
#include "coin2/app/driver.h"
#include "coin2/base/config/app_config.h"
#include "coin2/base/config/args.h"
#include "coin2/base/feed_convert_fastfeature.h"
#include "coin2/base/log.h"
#include "coin2/base/time.h"
#include "coin2/exchange/base/feed/update.h"
#include "coin2/strategy/linear_model/product_strategy.h"
#include "coin2/strategy/linear_model/util_ban_vu.h"
#include "coin2/strategy/strategy.h"
#include "presto/quant/base/init.h"
#include "presto/quant/base/module.h"
#include "presto/quant/feature/feature.h"
#include "presto/quant/feature/feature_dumper_papyrus.h"
#include "presto/quant/feature/feature_manager.h"
#include "presto/quant/symbology/symbol.h"
#include "presto/quant/symbology/symbol_manager.h"
#include "coin2/feed/fastfeed/system.h"

using namespace coin2::exchange::base::symbology;
using namespace coin2::base::config;
using namespace coin2::exchange::feed;
using namespace coin2::app;
using namespace coin2::strategy::linear_model;

// using ::presto::quant::feature::linear_model::LinearModel;
using ::coin2::base::ConvertFeedMessage;
using ::coin2::exchange::base::feed::FeedUpdate;
using ::coin2::strategy::linear_model::FeatureFeedConverter;
using ::presto::Symbol;
using ::presto::SymbolManager;
using ::presto::quant::feature::Dumper;
using ::util::json::DictionaryValue;


using coin2::feed::fastfeed::FastFeedUpdate;

// Helps a strategy to get imaginary clean feed.
struct ForwardLookHelper {
  typedef std::function<void(const FeedUpdate&)> callback_type;
  typedef std::tuple<int64_t, std::unique_ptr<FastFeedUpdate>, callback_type> entry;

  ForwardLookHelper(
      coin2::strategy::IStrategy* strategy,
      int64_t exchange_time_forward_look_sec)
      : strategy(strategy),
        exchange_time_forward_look_ns(exchange_time_forward_look_sec * 1e9) {}

  void UpdateFeed(
      int64_t exchange_time,
      const FeedUpdate& upd,
      callback_type callback) {
    if (exchange_time_forward_look_ns == 0L) {
      callback(upd);
    } else {
      if (exchange_time == 0) {
        exchange_time = upd.timestamp();
      }
      const FastFeedUpdate* upd_ptr = static_cast<const FastFeedUpdate*>(&upd);
      CHECK(upd_ptr);
      callback_queue.push(std::make_tuple(exchange_time, upd_ptr->Copy(), callback));
      UpdateTimeOnly(exchange_time);
    }
  }

  void UpdateTimeOnly(int64_t exchange_time) {
    while (!callback_queue.empty() &&
            std::get<0>(callback_queue.top()) < exchange_time - exchange_time_forward_look_ns) {
      const entry& entry_elem = callback_queue.top();
      const FastFeedUpdate* upd = std::get<1>(entry_elem).get();
      // unnecessary 1 more copy.. dislike this, but there's not too much of overhead I think.
      const_cast<FastFeedUpdate*>(upd)->SetTimestamp(std::get<0>(entry_elem));
      std::get<2>(entry_elem)(*upd);
      callback_queue.pop();
    }
  }

  static bool CompareEntries(const entry& a, const entry& b) {
    return std::get<0>(a) > std::get<0>(b);
  }

  coin2::strategy::IStrategy* strategy;
  int64_t exchange_time_forward_look_ns;
  std::priority_queue<entry, std::vector<entry>, decltype(&CompareEntries)> callback_queue{&CompareEntries};
};


class FastFeatureDumper : public coin2::strategy::IStrategy {
 public:
  std::string sampler_name_subset;
  std::experimental::filesystem::path feature_info_file;
  std::experimental::filesystem::path recorder_config_file;
  std::experimental::filesystem::path kline_json_file;
  std::string sample_from;
  std::string sample_until;
  int64_t timestamp_override;
  int64_t timestamp_min;

  FastFeatureDumper(
      const std::string& sampler_name_subset,
      const std::experimental::filesystem::path& feature_info_file,
      const std::experimental::filesystem::path& recorder_config_file,
      const std::experimental::filesystem::path& kline_json_file,
      const std::string& sample_from,
      const std::string& sample_until,
      int64_t timestamp_override,
      int64_t timestamp_min,
      int64_t exchange_time_forward_look_sec)
      : sampler_name_subset(sampler_name_subset),
        feature_info_file(feature_info_file),
        recorder_config_file(recorder_config_file),
        kline_json_file(kline_json_file),
        sample_from(sample_from),
        sample_until(sample_until),
        timestamp_override(timestamp_override),
        timestamp_min(timestamp_min),
        fwd_look_helper(std::make_unique<ForwardLookHelper>(this, exchange_time_forward_look_sec)) {}

  ~FastFeatureDumper() {
    // product might be hollow.. let's just give up.
    // fwd_look_helper->UpdateTimeOnly(timestamp_override + 3600000e9);
    LOG(INFO) << "Bye";
    dumper.reset();
  }

  void Init(::coin2::app::Driver* driver) override {
    std::vector<presto::Mep> meps;

    for (auto product : driver->feed()->GetProducts()) {
      meps.push_back(presto::Mep{
          MarketType_Name(product->market()),
          ExchangeType_Name(product->exchange()),
          product->relative_norm()
          });
    }

    LOG(INFO) << feature_info_file;
    LOG(INFO) << recorder_config_file;

    std::optional<int64_t> sample_from_ts;
    std::optional<int64_t> sample_until_ts;

    if (!sample_from.empty()) {
      sample_from_ts = TimestampFromString(sample_from);
    }

    if (!sample_until.empty()) {
      sample_until_ts = TimestampFromString(sample_until);
    }

    int64_t rank_ge = 0L;
    int64_t rank_lt = 0L;
    if (GetFocusSymbolsVuRankGeLt(
          GetFSpecProto(feature_info_file),
          &rank_ge,
          &rank_lt)) {
      std::string volume_info_filename = feature_info_file;
      boost::replace_all(volume_info_filename, "feature_info.json", "volume_info.json");
      coin2::strategy::linear_model::VolumeInfo volume_info;
      std::ifstream in(volume_info_filename);
      CHECK_THROW(in.is_open());
      std::string str((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
      CHECK(google::protobuf::util::JsonStringToMessage(str, &volume_info).ok())
          << volume_info_filename;
      auto feature_spec_pb = GetFSpecProto(feature_info_file);
      bool require_quote_matching = feature_spec_pb.focus_vu_require_quote_matching();
      bool exclude_noref_symbols = feature_spec_pb.focus_vu_exclude_noref_symbols();
      meps = BanFocusRankAndGetAllowlist(
          driver->feed(),
          volume_info,
          timestamp_override,
          rank_ge,
          rank_lt,
          nullptr,
          nullptr,
          require_quote_matching,
          exclude_noref_symbols);
    }

    std::string custom_data_filename = feature_info_file;
    boost::replace_all(custom_data_filename, "feature_info.json", "custom_data.json");
    std::ifstream in(custom_data_filename);
    std::unique_ptr<CustomData> custom_data;
    if (in.is_open()) {
      custom_data.reset(new CustomData());
      std::string str((std::istreambuf_iterator<char>(in)),
                      std::istreambuf_iterator<char>());
      google::protobuf::util::JsonStringToMessage(str, custom_data.get());
      presto::quant::feature::CutCustomData(custom_data.get(), timestamp_min);
    }

    dumper.reset(new Dumper(
        sampler_name_subset,
        feature_info_file,
        recorder_config_file,
        kline_json_file,
        meps,
        sample_from_ts,
        sample_until_ts,
        custom_data.get()));

    bool has_focus = dumper->GetFeatureManager()->IsFocus();
    converter.reset(new FeatureFeedConverter(
        dumper->GetFeatureManager()->GetSymbolManager(),
        &product_cache,
        true,
        has_focus ? std::nullopt : std::optional<int64_t>(2e9)));
  }

  void onFeatureFeed(const FeedUpdate& upd, const ::fastfeature::FeedMessage* feed_msg) {
    const Symbol* symbol = converter->GetSymbol(upd);
    // this feature dumper won't take this symbol.
    if (symbol == nullptr) return;
    dumper->UpdateFeed(symbol, *feed_msg);
  }

  void onTradeFeedImpl(const FeedUpdate& upd) {
    converter->onTradeFeed(upd);
    if (converter->feed_converted) {
      onFeatureFeed(upd, &converter->last_feed_msg);
    }
  }

  void onBookFeedImpl(const FeedUpdate& upd) {
    converter->onBookFeed(upd);
    if (converter->feed_converted) {
      onFeatureFeed(upd, &converter->last_feed_msg);
    }
  }

  void onStatusFeedImpl(const FeedUpdate& upd) {
    converter->onStatusFeed(upd);
    if (converter->feed_converted) {
      onFeatureFeed(upd, &converter->last_feed_msg);
    }
  }

  void onTradeFeed(const FeedUpdate& upd) override {
    auto callback = [this](const FeedUpdate&upd) { onTradeFeedImpl(upd); };
    fwd_look_helper->UpdateFeed(upd.trade().timestamp, upd, callback);
  }

  void onBookFeed(const FeedUpdate& upd) override {
    auto callback = [this](const FeedUpdate&upd) { onBookFeedImpl(upd); };
    fwd_look_helper->UpdateFeed(upd.book().Timestamp(), upd, callback);
  }

  void onStatusFeed(const FeedUpdate& upd) {
    auto callback = [this](const FeedUpdate&upd) { onStatusFeedImpl(upd); };
    int64_t exchange_time = 0L;
    if (upd.is_liquidation()) {
      exchange_time = upd.liquidation_order().timestamp;
    } else if (upd.is_index()) {
      exchange_time = upd.index().timestamp;
    } else if (upd.is_funding_rate()) {
      exchange_time = upd.funding_rate().timestamp;
    } else if (upd.is_open_interest()) {
      exchange_time = upd.open_interest().timestamp;
    } else if (upd.is_wallet()) {
      exchange_time = upd.wallet().timestamp;
    } else {
      // problematic, but not too problmatic frankly.
      exchange_time = upd.timestamp();
    }
    fwd_look_helper->UpdateFeed(exchange_time, upd, callback);
  }

  void onIndexFeed(const FeedUpdate& upd) override {
    onStatusFeed(upd);
  }

  void onLiquidationFeed(const FeedUpdate& upd) override {
    onStatusFeed(upd);
  }

  void onFundingRateFeed(const FeedUpdate& upd) override {
    onStatusFeed(upd);
  }

  void onOpenInterestFeed(const FeedUpdate& upd) override {
    onStatusFeed(upd);
  }

  void onWalletFeed(const FeedUpdate& upd) override {
    onStatusFeed(upd);
  }

 private:
  std::unique_ptr<Dumper> dumper;
  std::unique_ptr<FeatureFeedConverter> converter;
  ::coin2::exchange::base::symbology::ProductEncyclopedia product_cache;
  std::unique_ptr<ForwardLookHelper> fwd_look_helper;
};

int main(int argc, char* argv[]) {
  InitModules();

  DriverConfig app;

  cxxopts::Options opt("feature dumper", "feature dumper");
  opt.add_options()("sampler_name_subset", "", cxxopts::value<std::string>());
  opt.add_options()("config_filename", "feature config path", cxxopts::value<std::string>());
  opt.add_options()("recorder_config_filename", "", cxxopts::value<std::string>());
  opt.add_options()(
      "sample_from",
      "",
      cxxopts::value<std::string>()->default_value(""));
  opt.add_options()(
      "sample_until",
      "",
      cxxopts::value<std::string>()->default_value(""));
  opt.add_options()("kline_json", "kline info path",
      cxxopts::value<std::string>()->default_value(""));
  opt.add_options()("volume_info_filename", "volume info path",
      cxxopts::value<std::string>()->default_value(""));
  opt.add_options()("hours_warmedup", "hours_warmup that is already reflected",
      cxxopts::value<double>()->default_value("0.0"));
  opt.add_options()("exchange_time_forward_look_sec", "assume exchange timestamp comes on time",
      cxxopts::value<int64_t>()->default_value("0"));
  AddDriverOptions(&opt);
  AddIntervalOptions(&opt);
  AddProductOptions(&opt);
  AddFeedArchiveOptions(&opt);

  auto res = opt.parse(argc, argv);
  ParseDriverOptions(res, &app);
  ParseIntervalOptions(
      res,
      app.mutable_feed()->mutable_common()->mutable_archive()->mutable_interval());
  ParseProductOptions(res, app.mutable_feed());
  app.mutable_feed()->mutable_common()->mutable_archive()->set_machine(
      res["feed_machine"].as<std::string>());
  LOG(INFO) << app.DebugString();

  ::coin2::base::InitLogging(argv[0], app.log());

  auto begin_end = GetEnclosingInterval(app);
  CHECK_GT(begin_end.first, 0);

  int64_t nanosecs_warmedup = static_cast<int64_t>(res["hours_warmedup"].as<double>() * 3600 * 1e9);

  FastFeatureDumper fastfeature_dumper(
      res["sampler_name_subset"].as<std::string>(),
      res["config_filename"].as<std::string>(),
      res["recorder_config_filename"].as<std::string>(),
      res["kline_json"].as<std::string>(),
      res["sample_from"].as<std::string>(),
      res["sample_until"].as<std::string>(),
      begin_end.first + nanosecs_warmedup,
      begin_end.first,
      res["exchange_time_forward_look_sec"].as<int64_t>());
  {
    Driver driver(app, &fastfeature_dumper);
    driver.Run();
  }

  return 0;
}
